import cv2
import numpy as np
from typing import Optional
import time
from rknnpool import rknnPoolExecutor
import logging
logging.basicConfig(level=logging.WARNING)
from image_provider import ImageProvider
from logger import setup_logger
# 图像处理函数，实际应用过程中需要自行修改
from seg_inference2 import seg_func

logger = setup_logger()
# cap = cv2.VideoCapture(0)
segmodelPath = "./models/merge_yolov11_seg.rknn"
posemodelPath = "./models/merge_yolov11_pose.rknn"

class ImageData:
    """图像数据结构"""
    def __init__(self):
        self.timestamp: float = 0.0
        self.readable_timestamp: str = ""
        self.color_img: np.ndarray = None
        self.depth_img: np.ndarray = None
        self.color_path: Optional[str] = None
# 初始化线程池
# 线程数, 增大可提高帧率
TPEs = 3
# 初始化rknn池
seg_pool = rknnPoolExecutor(
    rknnModel=segmodelPath,
    TPEs=TPEs,
    func=seg_func,npus=[1,2])

# pose_pool = rknnPoolExecutor(
#     rknnModel=posemodelPath,
#     TPEs=TPEs,
#     func=seg_func,npus=[0])

###
image_provider = ImageProvider(from_folder= True)


frames, loopTime, initTime = 0, time.time(), time.time()
while True:
    image = image_provider.get_next_image()
    if image == None:
        break
    frames += 1
    frame = image.color_img
    seg_pool.put(frame)
    seg_result, flag = seg_pool.get()
    if flag == False:
        break
    print(seg_result)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break
    if frames % 30 == 0:
        print("30帧平均帧率:\t", 30 / (time.time() - loopTime), "帧")
        loopTime = time.time()